from scipy.optimize import linear_sum_assignment
from .cost import compute_cost_matrix
from .trajectory import generate_time_sequence


def assign_tasks(usvs, tasks, lambda_weight=0.1):
    cost_matrix = compute_cost_matrix(usvs, tasks, lambda_weight)
    row_ind, col_ind = linear_sum_assignment(cost_matrix)

    assignments = []
    for usv_idx, task_idx in zip(row_ind, col_ind):
        cost = cost_matrix[usv_idx][task_idx]
        if cost == float('inf'):
            continue
        usv = usvs[usv_idx]
        task = tasks[task_idx]
        trajectory = [usv['pos'], task['start_port'], task['end_port']]
        time_sequence = generate_time_sequence(trajectory, usv['speed'])
        assignments.append({
            'usv_id': usv['id'],
            'speed':usv['speed'],
            'task_id': task['id'],
            'payload': task['load'],
            'trajectory': trajectory,
            'time_sequence': time_sequence,
            'total_cost': cost
        })
    return assignments
